import cv2 as cv
import numpy as np
import sys

# load image
path = r'D:\_const\large_data\pic\classmates_requests\Li_Yujie_2022_05_19Thu.png'
img = cv.imread(path, cv.IMREAD_GRAYSCALE)
ret, img = cv.threshold(img, 1, 255, cv.THRESH_BINARY)
img_ = img.copy()
# cv.imshow('img', img_)

# find seed point for floodFill
h, w = img.shape
mask = np.zeros((h, w), dtype=np.uint8)
mask[:, w//2] = 255
# cv.imshow('mask', mask)
res = cv.bitwise_and(img, img, mask= mask)
# cv.imshow('res', res)
res_middle = res[:, w//2]
area1_y_min, area1_y_max, area2_y_min, area2_y_max = 0, 0, 0, 0
line1_over, line2_over, line3_over = -1, -1, -1
for i, val in enumerate(res_middle):
    if -1 == line1_over:
        if not val:
            continue
        else:
            line1_over = 0
            continue
    elif 0 == line1_over:
        if val:
            continue
        else:
            line1_over = 1
            area1_y_min = i
            continue
    elif -1 == line2_over:
        if not val:
            continue
        else:
            line2_over = 0
            area1_y_max = i - 1
            continue
    elif 0 == line2_over:
        if val:
            continue
        else:
            line2_over = 1
            area2_y_min = i
    elif -1 == line3_over:
        if not val:
            continue
        else:
            area2_y_max = i - 1
            break
print(area1_y_min, area1_y_max, area2_y_min, area2_y_max)
point1 = (w // 2, int(np.round((area1_y_min + area1_y_max) / 2)))
point2 = (w // 2, int(np.round((area2_y_min + area2_y_max) / 2)))
print(point1, point2)
cv.circle(img, point1, 1, 255, 1)
cv.circle(img, point2, 1, 255, 1)
cv.imshow('points', img)
img = img_.copy()

# floodFill
img3c = np.stack((img, img, img), axis=2)
# cv.imshow('img3c', img3c)
mask = np.zeros((h + 2, w + 2), dtype=np.uint8)
cv.floodFill(img3c, mask, point1, (0, 255, 0), 1, 1)
cv.imshow('img filled', img3c)

# split and calculate
unit_x = w // 5
mask_tpl = np.zeros_like(img, dtype=np.uint8)
area_unit = w * h
print(f'Area unit = {area_unit}')
for i in range(5):
    number = i + 1
    print(f'----{number}----')
    xleft = i * unit_x
    mask = mask_tpl.copy()
    mask[:, xleft:xleft+unit_x] = 255
    res = cv.bitwise_and(img3c, img3c, mask=mask)
    idx = (res == (255, 255, 255))
    # idx = (res == 255)  # The same effect. I guess numpy cannot compare a tuple but only can element by element, it is only broadcast.

    # study
    def check_idx(idx, name):
        print(f'--{name}--')
        print('idx.dtype', idx.dtype)
        print('idx.shape', idx.shape)
        print('idx unique', np.unique(idx))
        idx_mat = idx.astype(np.uint8) * 100
        cv.imshow(name, idx_mat)

    check_idx(idx, 'idx')
    check_idx(idx[:, :, 0], 'idx[:, :, 0]')
    check_idx(idx[:, :, 1], 'idx[:, :, 1]')
    check_idx(idx[:, :, 2], 'idx[:, :, 2]')

    # res[idx[:, :, 0]] = 0  # works
    # res[idx[:, :, 1]] = 0  # does not work, because floodFill color is (0, 255, 0)
    # res[idx[:, :, 2]] = 0  # works

    idx3c = np.all(idx, axis=2)  # ATTENTION check for 3 channels is available via this method
    # res[idx3c] = (0, 0, 0)
    res[idx3c] = (255, 0, 255)  # just for test

    # ret, res = cv.threshold(res[:, :, 1], 127, 255, cv.THRESH_BINARY)
    cv.imshow(f'#{number}', res)
    # area = cv.countNonZero(res)
    # print(f'Area = {area}, rate = {(area / area_unit):.4f}')

    cv.waitKey()
    cv.destroyAllWindows()
    sys.exit(0)
